import os
import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
import matplotlib.pyplot as plt

"""
silhouette_score Method for Optimal k
"""

max_label_num = 50000


def read_labels(label_file):
    objects = []
    with open(label_file, "r") as f:
        lines = f.readlines()
    for line in lines:
        parts = line.strip().split()
        if len(parts) > 1:
            width = float(parts[3]) - float(parts[1])
            height = float(parts[4]) - float(parts[2])
            objects.append((width, height))
    return objects


def calculate_aspect_ratio(objects):
    aspect_ratios = [w / h for w, h in objects if h != 0]
    return aspect_ratios


def plot_silhouette_score(aspect_ratios):
    silhouette_scores = []
    K_range = range(2, 6)
    for k in K_range:
        if len(aspect_ratios) < k:
            silhouette_scores.append(0)
            continue
        kmeans = KMeans(n_clusters=k, random_state=0)
        cluster_labels = kmeans.fit_predict(aspect_ratios)
        silhouette_avg = silhouette_score(aspect_ratios, cluster_labels)
        silhouette_avg = silhouette_score(aspect_ratios, cluster_labels)
        print(f"第{k}个：{silhouette_avg}")
        silhouette_scores.append(silhouette_avg)

    plt.figure(figsize=(10, 6))
    plt.plot(
        list(K_range), silhouette_scores, "bx-", marker="o", markersize=10, linewidth=3
    )

    for i, score in enumerate(silhouette_scores):
        if i == 0:
            plt.text(
                K_range[i] + 0.1,
                score,
                f"{score:.2f}",
                fontsize=12,
                ha="left",
                va="center",
            )
        elif i == 1:
            plt.text(
                K_range[i] + 0.1,
                score + 0.02,
                f"{score:.2f}",
                fontsize=12,
                ha="left",
                va="bottom",
            )
        else:
            plt.text(
                K_range[i],
                score + 0.02,
                f"{score:.2f}",
                fontsize=12,
                ha="center",
                va="bottom",
            )

    plt.xlabel("Number of Clusters (k)", fontsize=14)
    plt.ylabel("Silhouette Score", fontsize=14)
    plt.title("Silhouette Method for Optimal k", fontsize=16)
    plt.xticks(list(K_range), fontsize=12)
    plt.yticks(fontsize=12)
    plt.grid(True, linestyle="--", alpha=0.6)
    plt.tight_layout()

    # 保存图像为高分辨率的png文件
    plt.savefig("silhouette_score.png", dpi=1000, bbox_inches="tight", pad_inches=0.1)

    # 显示图像
    plt.show()


def main():
    folder_path = "/home/hw/dataset/cone_total_new/labels"
    label_files = [f for f in os.listdir(folder_path) if f.endswith(".txt")]

    aspect_ratios = []
    print("文件数量", len(label_files))
    print("计算中")
    for label_file in label_files:
        objects = read_labels(os.path.join(folder_path, label_file))
        aspect_ratios.extend(calculate_aspect_ratio(objects))

    aspect_ratios = np.array(aspect_ratios).reshape(-1, 1)
    print("目标数量", len(aspect_ratios))
    if len(aspect_ratios) > max_label_num:
        aspect_ratios = aspect_ratios[:max_label_num]
    print("最后目标数量", len(aspect_ratios))
    plot_silhouette_score(aspect_ratios)


if __name__ == "__main__":
    main()
